#include <iostream>
#include <stdio.h>
#include <cuda.h>
#include <curand.h>
#include "cuda_runtime.h"

#include "../time_cost.hpp"
#include "../gradient_compression_body/dgc_body.h"
#include "../operate_memory/get_policy_general.h"
#include "math.h"

#include <chrono>
#include <functional>
#include <fstream>
using namespace std;

int main(int argc, char** argv){
    int N = 10;
    double s_percent = 0.001;
    double  sample_rate = 0.001;
    int is_add_to = 0;
    int test_times = 100;
    int warmup_times = 10;
    if (argc>=2) N = atoi(argv[1]);
    if (argc>=3) s_percent = atof(argv[2]);
    if (argc>=4) sample_rate = atof(argv[3]);
    if (argc>=5) is_add_to = atoi(argv[4]);
    if (argc>=6) test_times = atoi(argv[5]);
    if (argc>=7) test_times = atoi(argv[6]);
    printf("N=%d\ts_percent=%lf\tsample_rate=%lf\t\n",
        N, s_percent, sample_rate);
    N = 1 << N;
    int M = 4*N;
    float* hostData;
    float* ORIGINAL;
    float* RESIDUAL;
    uint8_t* COMPRESSED;
    curandGenerator_t gen;
    cudaStream_t stream;
    cudaStreamCreate(&stream);
    hostData = (float*) calloc(N, sizeof(float));
    cudaMalloc((void**)&ORIGINAL, N*sizeof(float));
    cudaMalloc((void**)&RESIDUAL, N*sizeof(float));
    cudaMalloc((void**)&COMPRESSED, M*sizeof(uint8_t));
    curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);
    curandSetPseudoRandomGeneratorSeed(gen, 1234ULL);
    auto policy = zq_cpp_lib::operate_memory::get_policy<thrust::cuda_cub::par_t::stream_attachment_type>::get(stream);
    int ret;
    int total_compression_time = 0; // unit: microsecond(微秒)
    int total_decompression_time = 0;
    std::function<void(void)> init_data = [&](){
        curandGenerateUniform(gen, ORIGINAL, N);
        // cudaMemcpy(hostData, devData, N * sizeof(float), cudaMemcpyDeviceToHost);
    };
    std::function<void(void)> compress = [&](){
        auto start = chrono::system_clock::now();
        ret = zq_cpp_lib::gradient_compression_body::dgc_body<thrust::cuda_cub::par_t::stream_attachment_type>(
            ORIGINAL,
            N,
            COMPRESSED,
            sample_rate,
            s_percent,
            policy,
            stream
        );
        auto end = chrono::system_clock::now();
        auto duration = chrono::duration_cast<chrono::microseconds>(end-start);
        total_compression_time += duration.count();
    };
    std::function<void(void)> decompress = [&](){
        auto start = chrono::system_clock::now();
        ret = zq_cpp_lib::gradient_compression_body::dgc_r_body<thrust::cuda_cub::par_t::stream_attachment_type>(
            COMPRESSED,
            M,
            ORIGINAL,
            N,
            is_add_to,
            policy,
            stream
        );
        auto end = chrono::system_clock::now();
        auto duration = chrono::duration_cast<chrono::microseconds>(end-start);
        total_decompression_time += duration.count();
    };
    std::function<void(void)> run = [&](){
        init_data();
        compress();
        decompress();
    };
    for (int i = 0; i < warmup_times; i++){
        run();
    }
    total_compression_time = total_decompression_time = 0;
    for (int i = 0; i < test_times; i++){
        run();
    }
    
    double average_compression_cost = ((double)total_compression_time)/test_times;
    double average_decompression_cost = ((double)total_decompression_time)/test_times;
    printf("average compression cost %lf us\n", average_compression_cost);
    printf("average decompression cost %lf us\n", average_decompression_cost);
    ofstream out("ret.txt");
    out << average_compression_cost << endl;
    out << average_decompression_cost << endl;
    out.close();

    
    curandDestroyGenerator(gen);
    cudaFree(RESIDUAL);
    cudaFree(ORIGINAL);
    cudaFree(COMPRESSED);
    free(hostData);
    return 0;
}